[REFACTOR][IR] Use PrimType for compiler dtypes by tqchen · Pull Request #19875 · apache/tvm

tqchen · 2026-06-23T15:44:09Z

This PR removes the old runtime DataType wrapper as the compiler-facing dtype abstraction and routes dtype usage by boundary.

Rationale:

PrimType is the compiler IR type for primitive expressions, so using it as the dtype/type carrier unifies dtype information with Expr.ty.
This removes the need for a dedicated expression dtype path separate from type information.
Keeping dtype as an IR type leaves room for future expression type annotations without another parallel dtype abstraction.
Runtime, ABI, and dtype-valued attrs stay raw DLDataType, where a plain DLPack dtype value is the real boundary object.

Migration guide:

Use PrimType when code reasons about compiler expression types, tensor element compiler types, or constructs a PrimExpr/compiler type.
Use existing source types such as expr.ty(), ExprOp.expr_ty(), or TE tensor element dtype where possible instead of rebuilding a type from dtype text.
Use raw DLDataType for runtime constants, ABI paths, dtype-valued attrs, and storage/runtime helper logic.
Prefer direct PrimType equality, MatchesCode(...), MatchesElementType(...), and WithCode(...) over local wrappers or string dtype checks.
Keep public names such as GetDataType and output_dtype where they are API terminology, but align their value type with the compiler/runtime boundary.

Validation:

full branch changed-file pre-commit passed over 420 changed files
git diff --check passed
final grep found no PrimType::IsPredicate, .IsBool(, .IsInt(, .IsUInt(, removed type.cc helper names, or runtime/data_type.h references in checked paths
LLVM-enabled ninja -C build -j$(nproc) completed and linked lib/libtvm_compiler.so
focused Python pytest over TIRX ops/intrinsics, Relax qdq/manipulate legalizers, and contrib sort passed
direct TOPI smoke passed for touched math/scatter paths

The branch is intentionally stacked on #19874 (2bdedc93aa) so CI has the needed base fix.

gemini-code-assist

Code Review

This pull request refactors the codebase to transition from using DataType to PrimType (backed by DLDataType), unifying the type system across IR dialects and improving type safety. The reviewer provided valuable feedback on leveraging the new PrimType APIs (such as bits(), lanes(), and IsScalableVector()) to avoid manual DLDataType unpacking and arithmetic, particularly in tensor allocation, memory view handling, static memory planning, and GPU code verification. Additionally, a bug was identified in src/relax/transform/utils.h where kDLBool was incorrectly checked against a bit width of 1 instead of the standard 8 bits.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

tqchen · 2026-06-23T16:41:19Z

/gemini review

gemini-code-assist

Code Review

This pull request refactors the TVM codebase by replacing the usage of tvm::DataType with tvm::PrimType and DLDataType across C++ and Python files to better align with DLPack. The code review feedback is highly consistent, recommending the use of the more idiomatic MatchesCode helper on PrimType instead of directly comparing raw data type codes (such as kDLFloat, kDLInt, and kDLBool). Additionally, it is suggested to simplify redundant checks by passing multiple type codes to a single variadic MatchesCode call.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

tqchen · 2026-06-23T16:58:09Z

/gemini Please run another review pass on latest head b8c7505. Since the last pass, I applied the Gemini MatchesCode cleanups, fixed the Python PrimType-to-runtime-dtype boundary in Relax gradient constants and TE NumPy test uses, reran full LLVM build plus focused Relax/TE/TIRX validation, and resolved the previous review threads that are addressed.

gemini-code-assist

Code Review

This pull request refactors TVM's type system by replacing the runtime-specific DataType with a unified compile-time PrimType (backed by DLDataType) across the compiler, runtime, and Python bindings. This extensive refactoring updates buffer declarations, expression nodes, and codegen backends to use PrimType or DLDataType directly. Feedback on the changes highlights a potential division-by-zero bug in Hexagon's GetVectorBytes for sub-byte types, a potential AttributeError in the Python PrimExpr.dtype property when self.ty is null or a pointer, and several opportunities to simplify type checks using the newly introduced IsScalar() helper and PrimType equality operators.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-06-23T17:06:54Z

+  auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {
+    return to_type->dtype == expr.ty()->dtype ? expr : cast(to_type, expr);
  };


Since PrimType has operator== defined, the check to_type->dtype == expr.ty()->dtype can be simplified to to_type == expr.ty().

Suggested change

auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {

return to_type->dtype == expr.ty()->dtype ? expr : cast(to_type, expr);

};

auto cast_if_needed = [](PrimType to_type, PrimExpr expr) {

return to_type == expr.ty() ? expr : cast(to_type, expr);

};

gemini-code-assist · 2026-06-23T17:06:54Z

+  TVM_FFI_DCHECK(expr->BaseExprNode::ty.defined());
+  const auto* prim_ty = expr->BaseExprNode::ty.as<PrimTypeNode>();
+  TVM_FFI_DCHECK(prim_ty != nullptr);
+  return IsIndexType(prim_ty->dtype);


The manual checks on expr->BaseExprNode::ty are redundant because PrimExprNode::ty() already implements these checks and returns PrimType. We can simplify this to a single call to expr->ty().

Suggested change

TVM_FFI_DCHECK(expr->BaseExprNode::ty.defined());

const auto* prim_ty = expr->BaseExprNode::ty.as<PrimTypeNode>();

TVM_FFI_DCHECK(prim_ty != nullptr);

return IsIndexType(prim_ty->dtype);

return IsIndexType(expr->ty()->dtype);

gemini-code-assist · 2026-06-23T17:06:54Z


  bool IsValidScalar(const PrimExpr& expr) const {
-    return expr.defined() && expr.dtype().is_scalar();
+    if (!expr.defined()) return false;


The check !ty.IsFixedLengthVector() && !ty.IsScalableVector() can be simplified using the IsScalar() helper method on PrimType.

return expr.defined() && expr.ty().IsScalar();

tqchen · 2026-06-23T17:17:22Z

/gemini review

gemini-code-assist

Code Review

This pull request performs a major refactoring of TVM's type system by removing the tvm::runtime::DataType wrapper class and replacing its usage with tvm::PrimType and raw DLDataType across the entire codebase, including TIR, Relax, TOPI, and various hardware backends (CUDA, Vulkan, Metal, WebGPU, Hexagon, and Trainium). PrimExpr and PrimExprNode now expose their types via PrimType (ty()) instead of DataType (dtype()), and corresponding Python bindings have been updated to reflect these changes. Since no review comments were provided, there is no specific feedback to address, and the changes appear to successfully unify and simplify the type representation.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Use PrimType as the compiler-facing dtype/type carrier so primitive expression dtype information is unified with Expr.ty instead of flowing through a dedicated dtype path. This keeps compiler IR type information in the type system and leaves room for future expression type annotations. Use raw DLDataType at runtime, ABI, storage-helper, and dtype-valued attr boundaries where a plain DLPack dtype value is the real interface. Keep the PrimType API minimal and hot-path friendly with value equality, matching helpers, documented factories, and cached common constructors. Update TIRX, TE, TOPI, Relax, codegen, Python bindings, and tests to follow the compiler PrimType versus runtime DLDataType boundary.

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

tqchen force-pushed the task-use-ty-primtype-for-source-of-dtype branch 3 times, most recently from 3e75c55 to 74fba5e Compare June 23, 2026 16:32

gemini-code-assist Bot reviewed Jun 23, 2026

View reviewed changes

tqchen added 7 commits June 23, 2026 20:07

[REFACTOR][IR] Fix PrimType follow-up review issues

44abb6c

[REFACTOR][IR] Address PrimType review follow-ups

7e66589

[REFACTOR][Runtime] Use raw DLDataType in contrib backends

e130c82

[FIX][IR] Repair PrimType dtype CI fallout

7e3e9ce

[FIX][Vulkan] Avoid default PrimType construction

1313eb9

[FIX][Relax] Update reverse_sequence dtype checks

35d4a00

tqchen force-pushed the task-use-ty-primtype-for-source-of-dtype branch from b42ff43 to 35d4a00 Compare June 23, 2026 20:14

[FIX][Relax] Update CUTLASS dtype codegen

9789895

Conversation

tqchen commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tqchen commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tqchen commented Jun 23, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 23, 2026

Choose a reason for hiding this comment

Uh oh!

tqchen commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

tqchen commented Jun 23, 2026 •

edited

Loading

tqchen commented Jun 23, 2026 •

edited

Loading

tqchen commented Jun 23, 2026 •

edited

Loading